package com.dcits.tool.ssh;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.util.HashSet;
import java.util.Set;

import org.apache.log4j.Logger;

import ch.ethz.ssh2.Connection;
import ch.ethz.ssh2.Session;
import ch.ethz.ssh2.StreamGobbler;

public class SSHUtil {
	
	private static final Logger logger = Logger.getLogger(SSHUtil.class);
	
	private static final Set<String> EXEC_COMMAND_TAGS = new HashSet<String>();
	
	/**
	 * 获取ssh连接
	 * @param host
	 * @param port
	 * @param username
	 * @param password
	 * @return
	 * @throws IOException 
	 */
	public static Connection getConnection(String host, int port, String username, String password) throws IOException{
		Connection conn = new Connection(host, port);
		
		try {
			conn.connect(null, 3000, 2500);
			boolean flag = conn.authenticateWithPassword(username, password);
			if (flag) {
				return conn;
			}
		} catch (IOException e) {
			// TODO Auto-generated catch block
			logger.error(host + ":" + port + "连接失败!", e);		
			throw e;
		}	
		return null;
		
	}
	
	/**
	 * 发送停止标记
	 * @param tag
	 */
	public static void stopExecCommand(String tag) {
		EXEC_COMMAND_TAGS.add(tag);
	}
	
	/**
	 * 执行一次命令
	 * <br>返回结果
	 * <br>命令一定要是一次性执行完毕的命令(非交互性命令)
	 * 
	 * @param conn ssh连接对象
	 * @param command 执行命令
	 * @param count 读取多少行,从第一行开始读取
	 * @param getMode 返回模式，0-只返回正确的输出 1-只返回错误的信息 2-返回所有的信息<br>mode=2时将会使用终端模式
	 * @param tag 标记,防止中断命令时在并发情况下出现混乱
	 * @return
	 * @throws Exception 
	 * @throws IOException 
	 */
	public static String execCommand(Connection conn, String command, int count, int getMode, String tag) throws Exception {		
		StringBuilder str = new StringBuilder();
		
		if (conn != null) {
			Session session = null;
			InputStream is = null;
			BufferedReader brStat = null;
			try {
				session = conn.openSession(); 
								
				if (getMode == 1) {
					is = new StreamGobbler(session.getStderr());
				}
				
				if (getMode != 1) {
					is = new StreamGobbler(session.getStdout());
					if (getMode == 2) {
						session.requestPTY("vt100", 80, 24, 640, 480, null); 
					}
				}
												
				session.execCommand(command);
				brStat = new BufferedReader(new InputStreamReader(is));

				String readLine = null;
				int i = 0;
				while (i < count && (readLine = brStat.readLine()) != null) {					
					
					str.append(readLine);
					
					i++;
					
					if (i != count ) {
						str.append("\n");
					}
					//判断是否有中断命令
					if (EXEC_COMMAND_TAGS.contains(tag)) {
						break;
					}
				}				
			} catch (Exception e) {
				// TODO: handle exception
				e.printStackTrace();
				throw e;
			} finally {
				if (brStat != null) {
					try {
						brStat.close();
					} catch (IOException e) {
						// TODO Auto-generated catch block
						e.printStackTrace();
					}
				}
				if (is != null) {
					try {
						is.close();
					} catch (IOException e) {
						// TODO Auto-generated catch block
						e.printStackTrace();
					}	
				}
					
				if (session != null) {
					session.close();
				}
				
			}			
		}
		return str.toString();
	}
	
	/**
	 * 持续命令的执行
	 * <br>返回输出流
	 * @param conn
	 * @param command
	 * @return
	 * @throws IOException
	 */
	public static SSHBufferedReader execLoopCommand(Connection conn, String command) throws IOException {
		Session session = conn.openSession();

		InputStream is = new StreamGobbler(session.getStdout());		
		SSHBufferedReader brStat = new SSHBufferedReader(new InputStreamReader(is), session);
		session.execCommand(command);	
		return brStat;
	}
}
